#include "hypercube_hmmsq.hpp"
#include "intel_itt_wrapper.hpp"
#include <NTL/BasicThreadPool.h>

INTEL_ITT_DOMAIN_CREATE(hmmsq_compute_domain, "src.hypercube.hmmsq");
INTEL_ITT_STRING_HANDLE_CREATE(hmmsq_ra_handle, "RotateAlign");
INTEL_ITT_STRING_HANDLE_CREATE(hmmsq_sc_handle, "ShiftCompute");

static inline double log2_noise_bound(const helib::Ctxt &ctxt) {
  const double ln2 = NTL::log(NTL::xdouble(2l));
  return NTL::log(ctxt.totalNoiseBound()) / ln2;
}

void hypercube2d_rotate_align_ch::init(hmm_metric_collector *collector) {
  hmmsq_collector *hmmsq = static_cast<hmmsq_collector *>(collector);
  ASSERT_PRINTF(hmmsq != nullptr,
                "Failed to static_cast to `hmmsq_collector *`\n");

  const long h_nrow = ea_.sizeOfDimension(0);
  const long h_ncol = ea_.sizeOfDimension(1);
  const auto threads = std::min(threads_, static_cast<int>(h_nrow));

  //! Encoding in RotateAlign - Start
  lefts_.resize(h_nrow);
  rights_.resize(h_nrow);
  std::vector<NTL::zz_pX> lrow_i(h_nrow);
  std::vector<NTL::zz_pX> rrow_i(h_nrow);
#pragma omp parallel num_threads(threads)
  {
    auto tid = omp_get_thread_num();
#pragma omp for nowait
    for (long i = 0; i < h_nrow; i++) {
      hmmsq->start_encoding_ra_raw2vec(tid); //! Raw 2 Vec - Start
      std::vector<long> slots1(ea_.size(), 0);
      std::vector<long> slots2(ea_.size(), 0);
      for (long j = 0; j < h_ncol; j++) {
        slots1[i * h_ncol + j] = 1l;
      }
      for (long j = 0; j < h_nrow; j++) {
        slots2[j * h_ncol + i] = 1l;
      }
      hmmsq->stop_encoding_ra_raw2vec(tid); //! Raw 2 Vec - Completed

      hmmsq->start_encoding_ra_vec2ptxt(tid); //! Vec 2 Ptxt - Start
      ea_.encode(lefts_[i], slots1);
      ea_.encode(rights_[i], slots2);
      hmmsq->stop_encoding_ra_vec2ptxt(tid); //! Vec 2 Ptxt - Completed
    }

    //! Encoding in RotateAlign (Correct) - Start
    //!! `zz_pPush` must have a copy in each thread
    NTL::zz_pPush push(ea_.getPAlgebra().getP());
#pragma omp for
    for (long i = 0; i < h_nrow; i++) {
      hmmsq->start_encoding_ra_raw2vec(tid); //! Raw 2 Vec - Start
      std::vector<long> slots1(ea_.size(), 0);
      std::vector<long> slots2(ea_.size(), 0);
      for (long j = 0; j < h_ncol; j++) {
        if (i + j >= h_ncol) {
          slots1[i * h_ncol + j] = 1l;
        } else if (i + j < h_nrow) {
          slots2[i * h_ncol + j] = 1l;
        }
      }
      hmmsq->stop_encoding_ra_raw2vec(tid); //! Raw 2 Vec - Completed

      hmmsq->start_encoding_ra_vec2ptxt(tid); //! Vec 2 Ptxt - Start
      helib::zzX rrow, lrow;
      ea_.encode(rrow, slots1);
      ea_.encode(lrow, slots2);
      helib::convert(rrow_i[i], rrow);
      helib::convert(lrow_i[i], lrow);
      hmmsq->stop_encoding_ra_vec2ptxt(tid); //! Vec 2 Ptxt - Completed
    }
  }
  hmmsq->start_encoding_ra_re_ptxt(); //! Re/Ptxt - Start
  NTL::zz_pPush push(ea_.getPAlgebra().getP());
  NTL::zz_pX L, R;
  for (long i = 0; i < h_nrow; i++) {
    L += lrow_i[i];
    R += rrow_i[i];
  }
  correct_parts_[0] = helib::balanced_zzX(R);
  correct_parts_[1] = helib::balanced_zzX(L);
  hmmsq->stop_encoding_ra_re_ptxt(); //! Re/Ptxt - Start
  //! Encoding in RotateAlign (Correct) - Completed
  //! Encoding in RotateAlign - Completed
}

//* This version reuses encoded plaintexts in RotateAlign
void hypercube2d_shift_compute_ch::init(hmm_metric_collector *collector,
                                        const std::vector<helib::zzX> &cols) {
  hmmsq_collector *hmmsq = static_cast<hmmsq_collector *>(collector);
  ASSERT_PRINTF(hmmsq != nullptr,
                "Failed to static_cast to `hmmsq_collector *`\n");

  const long h_nrow = ea_.sizeOfDimension(0);
  const long h_ncol = ea_.sizeOfDimension(1);
  const auto threads = std::min(threads_, static_cast<int>(h_nrow));

  hmmsq->start_encoding_sc_re_ptxt(); //! Re/Ptxt - Start
  //* we compute on "encoded" plaintexts to avoid encoding full data
  correct_parts_.resize(h_nrow);
  std::vector<NTL::zz_pX> col_i(h_nrow);
#pragma omp parallel num_threads(threads)
  {
    //!! `zz_pPush` must have a copy in each thread
    NTL::zz_pPush push(ea_.getPAlgebra().getP());

    //* Step 1
    // Generate a constant polynomial corresponding to only column i as 1
#pragma omp for
    for (long i = 0; i < h_nrow; i++) {
      helib::convert(col_i[i], cols[i]);
    }

    //* Step 2
    // Using "after encoded" plaintexts to construct targets
    //* h_nrow^2 * \phi(m) times NTL::AddMod
#pragma omp for
    for (long i = 1; i < h_nrow; i++) {
      NTL::zz_pX L, R;
      for (long j = 0; j < h_nrow; j++) {
        j < h_nrow - i ? L += col_i[j] : R += col_i[j];
      }
      correct_parts_[i][0] = helib::balanced_zzX(L);
      correct_parts_[i][1] = helib::balanced_zzX(R);
    }
  }
  hmmsq->stop_encoding_sc_re_ptxt(); //! Re/Ptxt - Stop
}

/********************************************************************/
/*** Homomorphic Square Matrix Multiplication Based On Hypercube  ***/
/********************************************************************/

// Helper inline function for correcting RotateAlign
inline void correct_rotate_align_A(hmmsq_collector *hmmsq,
                                   hhmmsq_constant_multipler *hmmsqcm,
                                   const helib::EncryptedArray &ea,
                                   const long h_nrow, const long h_ncol,
                                   helib::Ctxt &A) {
  hmmsq->start_computing_ra_scale(); //! RA.Scale - Start
  helib::Ctxt tmp = A;
  hmmsqcm->rotate_align_correct_mask(tmp, A);
  hmmsq->stop_computing_ra_scale(); //! RA.Scale - Stop

  hmmsq->start_computing_ra_rot(); //! RA.Rot - Start
  ea.rotate1D(tmp, 1, -(h_ncol - h_nrow));
  hmmsq->stop_computing_ra_rot(); //! RA.Rot - Stop

  hmmsq->start_computing_ra_add(); //! RA.Add - Start
  A += tmp;
  hmmsq->stop_computing_ra_add(); //! RA.Add - Stop
}

// Helper inline function for ShiftCompute
inline void shift_compute(hmmsq_collector *hmmsq,
                          hhmmsq_constant_multipler *hmmsqcm,
                          const helib::EncryptedArray &ea, const long h_nrow,
                          const helib::Ctxt &Ap, const helib::Ctxt &Bp,
                          helib::Ctxt &AB) {
  // std::vector<helib::Ctxt> AB_partials(ea.sizeOfDimension(0),
  //                                      helib::Ctxt(AB.getPubKey()));
  hmmsq->start_computing_sc_rot(); //! SC.Rot - Start
  helib::GeneralAutomorphPrecon_FULL hoisted_Ap(Ap, 1, ea);
  helib::GeneralAutomorphPrecon_FULL hoisted_Bp(Bp, 0, ea);
  hmmsq->stop_computing_sc_rot(); //! SC.Rot - Stop
  for (int i = 0; i < h_nrow; i++) {
    hmmsq->start_computing_sc_rot(); //! SC.Rot - Start
    auto A_i = hoisted_Ap.automorph(-i);
    auto B_i = hoisted_Bp.automorph(-i);
    hmmsq->stop_computing_sc_rot(); //! SC.Rot - Stop

    // Correct Step: only 'A_i' required
    if (i) {
      hmmsq->start_computing_sc_rot(); //! SC.Rot - Start
      auto back = hoisted_Ap.automorph(h_nrow - i);
      hmmsq->stop_computing_sc_rot(); //! SC.Rot - Stop

      hmmsq->start_computing_sc_scale(); //! SC.Scale - Start
      hmmsqcm->shift_compute_correct_mask(*A_i, *back, i);
      hmmsq->stop_computing_sc_scale(); //! SC.Scale - Stop

      hmmsq->start_computing_sc_add(); //! SC.Add - Start
      *A_i += *back;
      hmmsq->stop_computing_sc_add(); //! SC.Add - Stop
    }

    hmmsq->start_computing_sc_mult(); //! SC.Mult - Start
    A_i->multiplyBy(*B_i);
    hmmsq->stop_computing_sc_mult(); //! SC.Mult - Stop

    // AB_partials[i] = *A_i;

    hmmsq->start_computing_sc_add(); //! SC.Add - Start
    AB += *A_i;
    hmmsq->stop_computing_sc_add(); //! SC.Add - Stop
  }

  // hmmsq->start_computing_sc_add(); //! SC.Add - Start
  // for (int stride = 1; 2 * stride <= AB_partials.size(); stride *= 2) {
  //   for (int i = 0; 2 * stride * i + stride < AB_partials.size(); i++) {
  //     AB_partials[2 * stride * i] += AB_partials[2 * stride * i + stride];
  //   }
  // }
  // AB = std::move(AB_partials[0]);
  // hmmsq->stop_computing_sc_add(); //! SC.Add - Stop
}

// Single Thread Version
//* Allow A = A * B
void hhmm_sq(hmm_metric_collector *collector, hmm_constant_multipler *cm,
             const helib::EncryptedArray &ea, helib::Ctxt &AB,
             const helib::Ctxt &A, const helib::Ctxt &B) {
  auto *hmmsq = static_cast<hmmsq_collector *>(collector);
  ASSERT_PRINTF(hmmsq != nullptr,
                "Failed to static_cast to `hmmsq_collector *`\n");
  auto *hmmsqcm = static_cast<hhmmsq_constant_multipler *>(cm);
  ASSERT_PRINTF(hmmsqcm != nullptr,
                "Failed to static_cast to `hhmmsq_constant_multipler *`\n");

  hmmsq->init_noise = std::max(log2_noise_bound(A), log2_noise_bound(B));

  const long h_nrow = ea.sizeOfDimension(0);
  const long h_ncol = ea.sizeOfDimension(1);

  const helib::PubKey &pk = A.getPubKey();
  helib::Ctxt zero_ctxt(pk);
  zero_ctxt.DummyEncrypt(NTL::ZZX(0));
  helib::Ctxt Ap(zero_ctxt);
  helib::Ctxt Bp(zero_ctxt);
  helib::Ctxt AB_temp(zero_ctxt);

  //! RotateAlign - Start
  hmmsq->start_computing_ra();
  hmmsq->start_computing_ra_rot(); //! RA.Rot - Start
  // Init rotation
  helib::GeneralAutomorphPrecon_FULL hoisted_A(A, 1, ea);
  helib::GeneralAutomorphPrecon_FULL hoisted_B(B, 0, ea);
  hmmsq->stop_computing_ra_rot(); //! RA.Rot - Stop

  for (long i = 0; i < h_nrow; i++) {
    hmmsq->start_computing_ra_rot(); //! RA.Rot - Start
    auto A_i = hoisted_A.automorph(-1l * i);
    auto B_i = hoisted_B.automorph(-1l * i);
    hmmsq->stop_computing_ra_rot(); //! RA.Rot - Stop

    hmmsq->start_computing_ra_scale(); //! RA.Scale - Start
    hmmsqcm->rotate_align_mask_left(*A_i, i);
    hmmsqcm->rotate_align_mask_right(*B_i, i);
    hmmsq->stop_computing_ra_scale(); //! RA.Scale - Stop

    hmmsq->start_computing_ra_add(); //! RA.Add - Start
    Ap += *A_i;
    Bp += *B_i;
    hmmsq->stop_computing_ra_add(); //! RA.Add - Stop
  }
  hmmsq->ra_noise =
      std::max(log2_noise_bound(Ap), log2_noise_bound(Bp)) - hmmsq->init_noise;

  //! Correct Ap (Correct RotateAlign) - Start
  hmmsq->ra_correct_noise = log2_noise_bound(Ap);
  correct_rotate_align_A(hmmsq, hmmsqcm, ea, h_nrow, h_ncol, Ap);
  hmmsq->ra_correct_noise = log2_noise_bound(Ap) - hmmsq->ra_correct_noise;
  //! Correct Ap (Correct RotateAlign) - Completed
  hmmsq->stop_computing_ra();
  //! RotateAlign - Completed

  //! ShiftCompute - Start
  hmmsq->sc_noise = std::max(log2_noise_bound(Ap), log2_noise_bound(Bp));
  hmmsq->start_computing_sc();
  shift_compute(hmmsq, hmmsqcm, ea, h_nrow, Ap, Bp, AB_temp);
  hmmsq->stop_computing_sc();
  //! ShiftCompute - Completed

  AB = std::move(AB_temp);
  hmmsq->sc_noise = log2_noise_bound(AB) - hmmsq->sc_noise;
  hmmsq->final_noise = log2_noise_bound(AB);
  hmmsq->left_cap = AB.capacity();
}

// Single Thread BSGS version
void hhmm_sq_bsgs(hmm_metric_collector *collector, hmm_constant_multipler *cm,
                  const helib::EncryptedArray &ea, helib::Ctxt &AB,
                  const helib::Ctxt &A, const helib::Ctxt &B) {
  auto *hmmsq = static_cast<hmmsq_collector *>(collector);
  ASSERT_PRINTF(hmmsq != nullptr,
                "Failed to static_cast to `hmmsq_collector *`\n");
  auto *hmmsqcm = static_cast<hhmmsq_constant_multipler *>(cm);
  ASSERT_PRINTF(hmmsqcm != nullptr,
                "Failed to static_cast to `hhmmsq_constant_multipler *`\n");

  hmmsq->init_noise = std::max(log2_noise_bound(A), log2_noise_bound(B));

  const long h_nrow = ea.sizeOfDimension(0);
  const long h_ncol = ea.sizeOfDimension(1);

  const long f = static_cast<long>(std::ceil(std::sqrt(h_nrow)));
  const long h = (h_nrow + f - 1) / f;

  const helib::PubKey &pk = A.getPubKey();
  helib::Ctxt zero_ctxt(pk);
  zero_ctxt.DummyEncrypt(NTL::ZZX(0));
  helib::Ctxt Ap(zero_ctxt);
  helib::Ctxt Bp(zero_ctxt);
  helib::Ctxt AB_temp(zero_ctxt);

  //! RotateAlign - Start
  hmmsq->start_computing_ra();
  std::vector<std::shared_ptr<helib::Ctxt>> A_baby(f);
  std::vector<std::shared_ptr<helib::Ctxt>> B_baby(f);

  hmmsq->start_computing_ra_rot(); //! RA.Rot - Start
  // Init rotation
  helib::GeneralAutomorphPrecon_FULL hoisted_A(A, 1, ea);
  helib::GeneralAutomorphPrecon_FULL hoisted_B(B, 0, ea);
  hmmsq->stop_computing_ra_rot(); //! RA.Rot - Stop

  // Pre-compute ctxts in the baby steps
  for (long i = 0; i < f; i++) {
    hmmsq->start_computing_ra_rot(); //! RA.Rot - Start
    A_baby[i] = hoisted_A.automorph(-1l * i);
    B_baby[i] = hoisted_B.automorph(-1l * i);
    hmmsq->stop_computing_ra_rot(); //! RA.Rot - Stop
  }

  // BSGS
  for (long i = 0; i < h; i++) {
    helib::Ctxt A_baby_sum(zero_ctxt);
    helib::Ctxt B_baby_sum(zero_ctxt);

    const long gaint_step = f * i;

    // Baby Steps
    for (long j = 0; j < f && gaint_step + j < h_nrow; j++) {
      hmmsq->start_computing_ra_rot(); //! RA.Rot - Start
      auto A_j = *A_baby[j];
      auto B_j = *B_baby[j];
      hmmsq->stop_computing_ra_rot(); //! RA.Rot - Stop

      hmmsq->start_computing_ra_scale(); //! RA.Scale - Start
      hmmsqcm->rotate_align_mask_left(A_j, j + gaint_step);
      hmmsqcm->rotate_align_mask_right(B_j, j + gaint_step);
      hmmsq->stop_computing_ra_scale(); //! RA.Scale - Stop

      hmmsq->start_computing_ra_add(); //! RA.Add - Start
      A_baby_sum += A_j;
      B_baby_sum += B_j;
      hmmsq->stop_computing_ra_add(); //! RA.Add - Stop
    }

    // Gaint Steps
    hmmsq->start_computing_ra_rot(); //! RA.Rot - Start
    ea.rotate1D(A_baby_sum, 1, -1l * gaint_step);
    ea.rotate1D(B_baby_sum, 0, -1l * gaint_step);
    hmmsq->stop_computing_ra_rot(); //! RA.Rot - Stop

    hmmsq->start_computing_ra_add(); //! RA.Add - Start
    Ap += A_baby_sum;
    Bp += B_baby_sum;
    hmmsq->stop_computing_ra_add(); //! RA.Add - Stop
  }
  hmmsq->ra_noise =
      std::max(log2_noise_bound(Ap), log2_noise_bound(Bp)) - hmmsq->init_noise;

  //! Correct Ap (Correct RotateAlign) - Start
  hmmsq->ra_correct_noise = log2_noise_bound(Ap);
  correct_rotate_align_A(hmmsq, hmmsqcm, ea, h_nrow, h_ncol, Ap);
  hmmsq->ra_correct_noise = log2_noise_bound(Ap) - hmmsq->ra_correct_noise;
  //! Correct Ap (Correct RotateAlign) - Completed
  hmmsq->stop_computing_ra();
  //! RotateAlign - Completed

  //! ShiftCompute - Start
  hmmsq->sc_noise = std::max(log2_noise_bound(Ap), log2_noise_bound(Bp));
  hmmsq->start_computing_sc();
  shift_compute(hmmsq, hmmsqcm, ea, h_nrow, Ap, Bp, AB_temp);
  hmmsq->stop_computing_sc();
  //! ShiftCompute - Completed

  AB = std::move(AB_temp);
  hmmsq->sc_noise = log2_noise_bound(AB) - hmmsq->sc_noise;
  hmmsq->final_noise = log2_noise_bound(AB);
  hmmsq->left_cap = AB.capacity();
}

// Multi-thread Version
void hhmm_sq_mt(hmm_metric_collector *collector, hmm_constant_multipler *cm,
                int threads, const helib::EncryptedArray &ea, helib::Ctxt &AB,
                const helib::Ctxt &A, const helib::Ctxt &B) {
  hmmsq_collector *hmmsq = static_cast<hmmsq_collector *>(collector);
  ASSERT_PRINTF(hmmsq != nullptr,
                "Failed to static_cast to `hmmsq_collector *`\n");
  auto *hmmsqcm = static_cast<hhmmsq_constant_multipler *>(cm);
  ASSERT_PRINTF(hmmsqcm != nullptr,
                "Failed to static_cast to `hhmmsq_constant_multipler *`\n");

  hmmsq->init_noise = std::max(log2_noise_bound(A), log2_noise_bound(B));

  const long h_nrow = ea.sizeOfDimension(0);
  const long h_ncol = ea.sizeOfDimension(1);

  const helib::PubKey &pk = A.getPubKey();
  helib::Ctxt zero_ctxt(pk);
  zero_ctxt.DummyEncrypt(NTL::ZZX(0));
  helib::Ctxt Ap(zero_ctxt);
  helib::Ctxt Bp(zero_ctxt);

  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_A;
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_B;
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_Ap;
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_Bp;
#ifndef OMP_CTXT_REDUCTION
  helib::Ctxt AB_temp(zero_ctxt);
#pragma omp parallel num_threads(threads)
  {
    INTEL_ITT_TASK_BEGIN(hmmsq_compute_domain, hmmsq_ra_handle); //! ITT - Begin
#pragma omp master               //! No implied barrier
    hmmsq->start_computing_ra(); //! RotateAlign - Start
#pragma omp single nowait
    hoisted_A = std::make_shared<helib::GeneralAutomorphPrecon_FULL>(A, 1, ea);
#pragma omp single nowait
    hoisted_B = std::make_shared<helib::GeneralAutomorphPrecon_FULL>(B, 0, ea);
#pragma omp barrier //! Explicit synchronization

    std::shared_ptr<helib::Ctxt> A_i;
    std::shared_ptr<helib::Ctxt> B_i;
#pragma omp for
    for (int i = 0; i < h_nrow; i++) {
      A_i = hoisted_A->automorph(-1 * i);
      B_i = hoisted_B->automorph(-1 * i);

      hmmsqcm->rotate_align_mask_left(*A_i, i);
      hmmsqcm->rotate_align_mask_right(*B_i, i);
#pragma omp critical
      {
        Ap += *A_i;
        Bp += *B_i;
      }
    }
    //! omp for: Implicit synchronization

#pragma omp single
    {
      hmmsq->ra_noise = std::max(log2_noise_bound(Ap), log2_noise_bound(Bp)) -
                        hmmsq->init_noise;

      //! Correct Ap (Correct RotateAlign) - Start
      hmmsq->ra_correct_noise = log2_noise_bound(Ap);
      helib::Ctxt Ap_temp = Ap;
      hmmsqcm->rotate_align_correct_mask(Ap_temp, Ap);
      ea.rotate1D(Ap_temp, 1l, -(h_ncol - h_nrow));
      Ap += Ap_temp;
      hmmsq->ra_correct_noise = log2_noise_bound(Ap) - hmmsq->ra_correct_noise;
      //! Correct Ap (Correct RotateAlign) - Completed
      hmmsq->sc_noise = std::max(log2_noise_bound(Ap), log2_noise_bound(Bp));
      hmmsq->stop_computing_ra(); //! RotateAlign - Completed
    }
    INTEL_ITT_TASK_END(hmmsq_compute_domain); //! ITT - End

    INTEL_ITT_TASK_BEGIN(hmmsq_compute_domain, hmmsq_sc_handle); //! ITT - Begin
#pragma omp single nowait
    {
      hmmsq->start_computing_sc(); //! ShiftCompute - Start
      hoisted_Ap =
          std::make_shared<helib::GeneralAutomorphPrecon_FULL>(Ap, 1, ea);
    }
#pragma omp single nowait
    hoisted_Bp =
        std::make_shared<helib::GeneralAutomorphPrecon_FULL>(Bp, 0, ea);
#pragma omp barrier //! Explicit synchronization

#pragma omp for
    for (int i = 0; i < h_nrow; i++) {
      // Shift
      A_i = hoisted_Ap->automorph(-1l * i);
      B_i = hoisted_Bp->automorph(-1l * i);

      // Correct Step: only 'A_i'(i>0) required
      if (i) {
        auto back = hoisted_Ap->automorph(h_nrow - i);
        hmmsqcm->shift_compute_correct_mask(*A_i, *back, i);
        *A_i += *back;
      }
      A_i->multiplyBy(*B_i);

#pragma omp critical
      { AB_temp += *A_i; }
    }
    //! omp for: Implicit synchronization

#pragma omp master                            //! No implied barrier
    hmmsq->stop_computing_sc();               //! ShiftCompute - Completed
    INTEL_ITT_TASK_END(hmmsq_compute_domain); //! ITT - End
  }
#else
  helib::Ctxt AB_temp(zero_ctxt);
  std::vector<std::shared_ptr<helib::Ctxt>> A_acc(h_nrow);
  std::vector<std::shared_ptr<helib::Ctxt>> B_acc(h_nrow);
#pragma omp parallel num_threads(threads)
  {
#pragma omp master //! No implied barrier
    hmmsq->start_computing_ra(); //! RotateAlign - Start
#pragma omp single nowait
    hoisted_A = std::make_shared<helib::GeneralAutomorphPrecon_FULL>(A, 1, ea);
#pragma omp single nowait
    hoisted_B = std::make_shared<helib::GeneralAutomorphPrecon_FULL>(B, 0, ea);
#pragma omp barrier //! Explicit synchronization

#pragma omp for
    for (int i = 0; i < h_nrow; i++) {
      A_acc[i] = hoisted_A->automorph(-1l * i);
      B_acc[i] = hoisted_B->automorph(-1l * i);

      hmmsqcm->rotate_align_mask_left(*A_acc[i], i);
      hmmsqcm->rotate_align_mask_right(*B_acc[i], i);
    }

#pragma omp for reduction(ctxt_plus : Ap, Bp)
    for (int i = 0; i < h_nrow; i++) {
#pragma omp critical
      Ap += *A_acc[i];
      Bp += *B_acc[i];
    }
    //! omp for: Implicit synchronization

    //? load balance
    //* CORRECT -+--> ATM_PRE_COMP --------> A_acc[i] * B_acc[i] (i>=1)
    //*          +--> A_acc[0] * B_acc[0]
#pragma omp single nowait
    {
#pragma omp task depend(out : Ap)
      {
        //! Correct Ap (Correct RotateAlign) - Start
        helib::Ctxt Ap_temp = Ap;
        hmmsqcm->rotate_align_correct_mask(Ap_temp, Ap);
        ea.rotate1D(Ap_temp, 1, -(h_ncol - h_nrow));
        Ap += Ap_temp;
        //! Correct Ap (Correct RotateAlign) - Completed
        hmmsq->stop_computing_ra(); //! RotateAlign - Completed
      }
#pragma omp task depend(in : Ap)
      {
        hmmsq->start_computing_sc(); //! ShiftCompute - Start
        hoisted_Ap =
            std::make_shared<helib::GeneralAutomorphPrecon_FULL>(Ap, 1, ea);
      }
#pragma omp task depend(in : Ap)
      {
        A_acc[0] = std::make_shared<helib::Ctxt>(Ap);
        A_acc[0]->multiplyBy(Bp);
      }
    }
#pragma omp single nowait
    hoisted_Bp =
        std::make_shared<helib::GeneralAutomorphPrecon_FULL>(Bp, 0, ea);
#pragma omp barrier //! Explicit synchronization

#pragma omp for
    for (int i = 1; i < h_nrow; i++) {
      A_acc[i] = hoisted_Ap->automorph(-1l * i);
      B_acc[i] = hoisted_Bp->automorph(-1l * i);

      // Correct Step: only 'A_i' required
      auto back = hoisted_Ap->automorph(h_nrow - i);
      hmmsqcm->shift_compute_correct_mask(*A_acc[i], *back, i);
      *A_acc[i] += *back;

      A_acc[i]->multiplyBy(*B_acc[i]);
    }
    //! omp for: Implicit synchronization
    //* So we do not need to worry about the data-race here

    // Reduction
#pragma omp for reduction(ctxt_plus : AB_temp)
    for (int i = 0; i < h_nrow; i++) {
      AB_temp += *A_acc[i];
    }
#pragma omp master //! No implied barrier
    hmmsq->stop_computing_sc(); //! ShiftCompute - Completed
  }
#endif

  AB = std::move(AB_temp);
  hmmsq->sc_noise = log2_noise_bound(AB) - hmmsq->sc_noise;
  hmmsq->final_noise = log2_noise_bound(AB);
  hmmsq->left_cap = AB.capacity();
}

void hhmm_sq_bsgs_mt_task(hmmsq_collector *hmmsq,
                          hhmmsq_constant_multipler *hmmsqcm,
                          const helib::EncryptedArray &ea, helib::Ctxt &AB,
                          const helib::Ctxt &A, const helib::Ctxt &B) {
  hmmsq->init_noise = std::max(log2_noise_bound(A), log2_noise_bound(B));

  const long h_nrow = ea.sizeOfDimension(0);
  const long h_ncol = ea.sizeOfDimension(1);

  const long f = static_cast<long>(std::ceil(std::sqrt(h_nrow)));
  const long h = (h_nrow + f - 1) / f;

  const helib::PubKey &pk = A.getPubKey();
  helib::Ctxt zero_ctxt(pk);
  zero_ctxt.DummyEncrypt(NTL::ZZX(0));
  helib::Ctxt Ap(zero_ctxt);
  helib::Ctxt Bp(zero_ctxt);
  helib::Ctxt AB_temp(zero_ctxt);
  // std::vector<std::shared_ptr<helib::Ctxt>> C(h_nrow);
  // std::vector<std::shared_ptr<helib::Ctxt>> Cp(h_nrow);

  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_A;
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_B;

  //! RotateAlign - Start
  std::vector<std::shared_ptr<helib::Ctxt>> A_baby(f);
  std::vector<std::shared_ptr<helib::Ctxt>> B_baby(f);
  std::vector<helib::Ctxt> A_baby_sum(h, zero_ctxt);
  std::vector<helib::Ctxt> B_baby_sum(h, zero_ctxt);
#pragma omp parallel
#pragma omp single
  {
    INTEL_ITT_TASK_BEGIN(hmmsq_compute_domain, hmmsq_ra_handle); //! ITT - Begin
    hmmsq->start_computing_ra(); //! RotateAlign - Start

#pragma omp task
    hoisted_A = std::make_shared<helib::GeneralAutomorphPrecon_FULL>(A, 1, ea);
#pragma omp task
    hoisted_B = std::make_shared<helib::GeneralAutomorphPrecon_FULL>(B, 0, ea);
#pragma omp taskwait //! Explicit synchronization

    // Pre-compute ctxts in the baby steps
#pragma omp taskloop
    for (long i = 0; i < f; i++) {
      A_baby[i] = hoisted_A->automorph(-1l * i);
      B_baby[i] = hoisted_B->automorph(-1l * i);
    } //! omp taskloop: Implicit synchronization

    // BSGS
    // FIXME: more parallelism?
#pragma omp taskloop
    for (long i = 0; i < h; i++) {
      // helib::Ctxt A_baby_sum(zero_ctxt);
      // helib::Ctxt B_baby_sum(zero_ctxt);

      const long gaint_step = f * i;

      // Baby Steps
      for (long j = 0; j < f && gaint_step + j < h_nrow; j++) {
        auto A_j = *A_baby[j];
        auto B_j = *B_baby[j];

        hmmsqcm->rotate_align_mask_left(A_j, j + gaint_step);
        hmmsqcm->rotate_align_mask_right(B_j, j + gaint_step);

        A_baby_sum[i] += A_j;
        B_baby_sum[i] += B_j;
      }

      // Gaint Steps
      ea.rotate1D(A_baby_sum[i], 1, -1l * gaint_step);
      ea.rotate1D(B_baby_sum[i], 0, -1l * gaint_step);

      // #pragma omp critical
      //       {
      //         Ap += A_baby_sum;
      //         Bp += B_baby_sum;
      //       }
    } //! omp taskloop: Implicit synchronization

#pragma omp task depend(out : Ap)
    {
      for (int i = 0; i < h; i++) {
        Ap += A_baby_sum[i];
      }
      hmmsq->ra_noise = std::max(log2_noise_bound(Ap), log2_noise_bound(Bp)) -
                        hmmsq->init_noise;

      //! Correct Ap (Correct RotateAlign) - Start
      hmmsq->ra_correct_noise = log2_noise_bound(Ap);
      helib::Ctxt Ap_temp = Ap;
      hmmsqcm->rotate_align_correct_mask(Ap_temp, Ap);
      ea.rotate1D(Ap_temp, 1, -(h_ncol - h_nrow));
      Ap += Ap_temp;
      hmmsq->ra_correct_noise = log2_noise_bound(Ap) - hmmsq->ra_correct_noise;
      //! Correct Ap (Correct RotateAlign) - Completed
      hmmsq->stop_computing_ra(); //! RotateAlign - Completed
      hmmsq->sc_noise = std::max(log2_noise_bound(Ap), log2_noise_bound(Bp));
    }
#pragma omp task depend(out : Bp)
    {
      for (int i = 0; i < h; i++) {
        Bp += B_baby_sum[i];
      }
    }
    INTEL_ITT_TASK_END(hmmsq_compute_domain); //! ITT - End

    INTEL_ITT_TASK_BEGIN(hmmsq_compute_domain, hmmsq_sc_handle); //! ITT - Begin
    hmmsq->start_computing_sc(); //! ShiftCompute - Start
#pragma omp task depend(in : Ap)
    hoisted_A = std::make_shared<helib::GeneralAutomorphPrecon_FULL>(Ap, 1, ea);
#pragma omp task depend(in : Bp)
    hoisted_B = std::make_shared<helib::GeneralAutomorphPrecon_FULL>(Bp, 0, ea);
#pragma omp taskwait //! Explicit synchronization

#pragma omp taskloop
    for (int i = 0; i < h_nrow; i++) {
      // Shift
      auto A_i = hoisted_A->automorph(-i);
      auto B_i = hoisted_B->automorph(-i);
      // Cp[i] = hoisted_A->automorph(-i);

      // Correct Step: only 'A_i'(i>0) required
      if (i) {
        auto back = hoisted_A->automorph(h_nrow - i);
        hmmsqcm->shift_compute_correct_mask(*A_i, *back, i);
        *A_i += *back;
        // hmmsqcm->shift_compute_correct_mask(*Cp[i], *back, i);
        // *Cp[i] += *back;
      }
      A_i->multiplyBy(*B_i);

#pragma omp critical
      { AB_temp += *A_i; }
    } //! omp taskloop: Implicit synchronization

    // #pragma omp taskloop
    // for (int i = 0; i < h_nrow; i++) {
    // C[i] = hoisted_B->automorph(-i);
    // C[i]->multiplyBy(*Cp[i]);
    // } //! omp taskloop: Implicit synchronization

    hmmsq->stop_computing_sc();               //! ShiftCompute - Completed
    INTEL_ITT_TASK_END(hmmsq_compute_domain); //! ITT - End
  }
  // AB = std::move(zero_ctxt);
  // // NTL::SetNumThreads(omp_get_max_threads());
  // for (auto i = 0; i < C.size(); i++) {
  //   AB += *C[i];
  // }
  // NTL::SetNumThreads(1l);
  AB = std::move(AB_temp);
  hmmsq->sc_noise = log2_noise_bound(AB) - hmmsq->sc_noise;
  hmmsq->final_noise = log2_noise_bound(AB);
  hmmsq->left_cap = AB.capacity();
}

void hhmm_sq_rotate_align_A_batch(hmm_constant_multipler *cm,
                                  const helib::EncryptedArray &ea,
                                  int rot_pos_A, int rot_size_A,
                                  helib::Ctxt &Ap, const helib::Ctxt &A) {
  auto *hmmsqcm = static_cast<hhmmsq_constant_multipler *>(cm);
  ASSERT_PRINTF(hmmsqcm != nullptr,
                "Failed to static_cast to `hhmmsq_constant_multipler *`\n");

  const long h_nrow = ea.sizeOfDimension(0);
  const long h_ncol = ea.sizeOfDimension(1);
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_A =
      std::make_shared<helib::GeneralAutomorphPrecon_FULL>(A, 1, ea);

  if (rot_size_A != h_nrow) {
    std::vector<std::shared_ptr<helib::Ctxt>> Ais(rot_size_A);
#pragma omp parallel
#pragma omp single
#pragma omp taskloop
    for (int i = 0; i < rot_size_A; i++) {
      long r = rot_pos_A + i;
      Ais[i] = hoisted_A->automorph(-1l * r);
      hmmsqcm->rotate_align_mask_left(*Ais[i], r);
    }
    for (int i = 1; i < Ais.size(); i++) {
      Ais[0]->addCtxt(*Ais[i]);
    }
    //! Correct Ap (Correct RotateAlign)
    helib::Ctxt Ap_tmp = *Ais[0];
    hmmsqcm->rotate_align_correct_mask(Ap_tmp, *Ais[0]);
    ea.rotate1D(Ap_tmp, 1, -(h_ncol - h_nrow));
    Ais[0]->addCtxt(Ap_tmp);

    Ap = std::move(*Ais[0]);
  } else {
    helib::Ctxt zero_ctxt(A.getPubKey());
    zero_ctxt.DummyEncrypt(NTL::ZZX(0));
    const long f = static_cast<long>(std::ceil(std::sqrt(h_nrow)));
    const long h = (h_nrow + f - 1) / f;
    std::vector<std::shared_ptr<helib::Ctxt>> A_baby(f);
    std::vector<helib::Ctxt> A_baby_sum(h, zero_ctxt);

#pragma omp parallel
#pragma omp single
    {
      // Pre-compute ctxts in the baby steps
#pragma omp taskloop
      for (long i = 0; i < f; i++) {
        A_baby[i] = hoisted_A->automorph(-1l * i);
      }
      // BSGS
#pragma omp taskloop
      for (long i = 0; i < h; i++) {
        const long gaint_step = f * i;
        // Baby Steps
        for (long j = 0; j < f && gaint_step + j < h_nrow; j++) {
          auto A_j = *A_baby[j];
          hmmsqcm->rotate_align_mask_left(A_j, j + gaint_step);
          A_baby_sum[i] += A_j;
        }
        // Gaint Steps
        ea.rotate1D(A_baby_sum[i], 1, -1l * gaint_step);
      }
    }
    for (int i = 1; i < h; i++) {
      A_baby_sum[0] += A_baby_sum[i];
    }
    //! Correct Ap (Correct RotateAlign)
    helib::Ctxt Ap_tmp = A_baby_sum[0];
    hmmsqcm->rotate_align_correct_mask(Ap_tmp, A_baby_sum[0]);
    ea.rotate1D(Ap_tmp, 1, -(h_ncol - h_nrow));
    A_baby_sum[0].addCtxt(Ap_tmp);

    Ap = std::move(A_baby_sum[0]);
  }
}

void hhmm_sq_rotate_align_B_batch(hmm_constant_multipler *cm,
                                  const helib::EncryptedArray &ea,
                                  int rot_pos_B, int rot_size_B,
                                  helib::Ctxt &Bp, const helib::Ctxt &B) {
  auto *hmmsqcm = static_cast<hhmmsq_constant_multipler *>(cm);
  ASSERT_PRINTF(hmmsqcm != nullptr,
                "Failed to static_cast to `hhmmsq_constant_multipler *`\n");

  const long h_nrow = ea.sizeOfDimension(0);
  const long h_ncol = ea.sizeOfDimension(1);
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_B =
      std::make_shared<helib::GeneralAutomorphPrecon_FULL>(B, 0, ea);

  if (rot_size_B != h_nrow) {
    std::vector<std::shared_ptr<helib::Ctxt>> Bis(rot_size_B);
#pragma omp parallel
#pragma omp single
#pragma omp taskloop
    for (int i = 0; i < rot_size_B; i++) {
      long r = rot_pos_B + i;
      Bis[i] = hoisted_B->automorph(-1l * r);
      hmmsqcm->rotate_align_mask_right(*Bis[i], r);
    }
    for (int i = 1; i < Bis.size(); i++) {
      Bis[0]->addCtxt(*Bis[i]);
    }
    Bp = std::move(*Bis[0]);
  } else {
    helib::Ctxt zero_ctxt(B.getPubKey());
    zero_ctxt.DummyEncrypt(NTL::ZZX(0));
    const long f = static_cast<long>(std::ceil(std::sqrt(h_nrow)));
    const long h = (h_nrow + f - 1) / f;
    std::vector<std::shared_ptr<helib::Ctxt>> B_baby(f);
    std::vector<helib::Ctxt> B_baby_sum(h, zero_ctxt);

#pragma omp parallel
#pragma omp single
    {
      // Pre-compute ctxts in the baby steps
#pragma omp taskloop
      for (long i = 0; i < f; i++) {
        B_baby[i] = hoisted_B->automorph(-1l * i);
      }
      // BSGS
#pragma omp taskloop
      for (long i = 0; i < h; i++) {
        const long gaint_step = f * i;
        // Baby Steps
        for (long j = 0; j < f && gaint_step + j < h_nrow; j++) {
          auto B_j = *B_baby[j];
          hmmsqcm->rotate_align_mask_right(B_j, j + gaint_step);
          B_baby_sum[i] += B_j;
        }
        // Gaint Steps
        ea.rotate1D(B_baby_sum[i], 0, -1l * gaint_step);
      }
    }
    for (int i = 1; i < h; i++) {
      B_baby_sum[0] += B_baby_sum[i];
    }
    Bp = std::move(B_baby_sum[0]);
  }
}

void hhmm_sq_shift_compute_A_batch(hmm_constant_multipler *cm,
                                   const helib::EncryptedArray &ea,
                                   int rot_pos_A, std::vector<helib::Ctxt> &Ais,
                                   const helib::Ctxt &A) {
  auto *hmmsqcm = static_cast<hhmmsq_constant_multipler *>(cm);
  ASSERT_PRINTF(hmmsqcm != nullptr,
                "Failed to static_cast to `hhmmsq_constant_multipler *`\n");

  const long h_nrow = ea.sizeOfDimension(0);
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_A =
      std::make_shared<helib::GeneralAutomorphPrecon_FULL>(A, 1, ea);
#pragma omp parallel
#pragma omp single
#pragma omp taskloop
  for (int i = 0; i < Ais.size(); i++) {
    long r = rot_pos_A + i;
    auto Ai = hoisted_A->automorph(-1l * r);
    if (r) {
      auto back = hoisted_A->automorph(h_nrow - r);
      hmmsqcm->shift_compute_correct_mask(*Ai, *back, r);
      *Ai += *back;
    }
    Ais[i] = *Ai;
  }
}

void hhmm_sq_shift_compute_B_batch(hmm_constant_multipler *cm,
                                   const helib::EncryptedArray &ea,
                                   int rot_pos_B, std::vector<helib::Ctxt> &Bis,
                                   const helib::Ctxt &B) {
  auto *hmmsqcm = static_cast<hhmmsq_constant_multipler *>(cm);
  ASSERT_PRINTF(hmmsqcm != nullptr,
                "Failed to static_cast to `hhmmsq_constant_multipler *`\n");

  const long h_nrow = ea.sizeOfDimension(0);
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_B =
      std::make_shared<helib::GeneralAutomorphPrecon_FULL>(B, 0, ea);
#pragma omp parallel
#pragma omp single
#pragma omp taskloop
  for (int i = 0; i < Bis.size(); i++) {
    long r = rot_pos_B + i;
    auto Bi = hoisted_B->automorph(-1l * r);
    Bis[i] = *Bi;
  }
}

void hhmm_sq_rotate_align_A(hmm_constant_multipler *cm,
                            const helib::EncryptedArray &ea, helib::Ctxt &Ap,
                            const helib::Ctxt &A) {
  auto *hmmsqcm = static_cast<hhmmsq_constant_multipler *>(cm);
  ASSERT_PRINTF(hmmsqcm != nullptr,
                "Failed to static_cast to `hhmmsq_constant_multipler *`\n");

  const long h_nrow = ea.sizeOfDimension(0);
  const long h_ncol = ea.sizeOfDimension(1);

  const long f = static_cast<long>(std::ceil(std::sqrt(h_nrow)));
  const long h = (h_nrow + f - 1) / f;

  helib::Ctxt zero_ctxt(A.getPubKey());
  zero_ctxt.DummyEncrypt(NTL::ZZX(0l));
  helib::Ctxt tmp(zero_ctxt);

  std::vector<std::shared_ptr<helib::Ctxt>> A_baby(f);

  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_A =
      std::make_shared<helib::GeneralAutomorphPrecon_FULL>(A, 1, ea);
  // Pre-compute ctxts in the baby steps
  for (long i = 0; i < f; i++) {
    A_baby[i] = hoisted_A->automorph(-1l * i);
  }
  // BSGS
  for (long i = 0; i < h; i++) {
    helib::Ctxt A_baby_sum(zero_ctxt);
    const long gaint_step = f * i;
    // Baby Steps
    for (long j = 0; j < f && gaint_step + j < h_nrow; j++) {
      auto A_j = *A_baby[j];
      hmmsqcm->rotate_align_mask_left(A_j, j + gaint_step);
      A_baby_sum += A_j;
    }
    // Gaint Steps
    ea.rotate1D(A_baby_sum, 1, -1l * gaint_step);
    tmp += A_baby_sum;
  }

  //! Correct Ap (Correct RotateAlign)
  helib::Ctxt Ap_tmp = tmp;
  hmmsqcm->rotate_align_correct_mask(Ap_tmp, tmp);
  ea.rotate1D(Ap_tmp, 1, -(h_ncol - h_nrow));
  tmp += Ap_tmp;

  Ap = std::move(tmp);
}

void hhmm_sq_rotate_align_B(hmm_constant_multipler *cm,
                            const helib::EncryptedArray &ea, helib::Ctxt &Bp,
                            const helib::Ctxt &B) {
  auto *hmmsqcm = static_cast<hhmmsq_constant_multipler *>(cm);
  ASSERT_PRINTF(hmmsqcm != nullptr,
                "Failed to static_cast to `hhmmsq_constant_multipler *`\n");

  const long h_nrow = ea.sizeOfDimension(0);
  const long h_ncol = ea.sizeOfDimension(1);

  const long f = static_cast<long>(std::ceil(std::sqrt(h_nrow)));
  const long h = (h_nrow + f - 1) / f;

  helib::Ctxt zero_ctxt(B.getPubKey());
  zero_ctxt.DummyEncrypt(NTL::ZZX(0l));
  helib::Ctxt tmp(zero_ctxt);

  std::vector<std::shared_ptr<helib::Ctxt>> B_baby(f);

  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_B =
      std::make_shared<helib::GeneralAutomorphPrecon_FULL>(B, 0, ea);
  // Pre-compute ctxts in the baby steps
  for (long i = 0; i < f; i++) {
    B_baby[i] = hoisted_B->automorph(-1l * i);
  }

  // BSGS
  for (long i = 0; i < h; i++) {
    helib::Ctxt B_baby_sum(zero_ctxt);
    const long gaint_step = f * i;
    // Baby Steps
    for (long j = 0; j < f && gaint_step + j < h_nrow; j++) {
      auto B_j = *B_baby[j];
      hmmsqcm->rotate_align_mask_right(B_j, j + gaint_step);
      B_baby_sum += B_j;
    }
    // Gaint Steps
    ea.rotate1D(B_baby_sum, 0, -1l * gaint_step);
    tmp += B_baby_sum;
  }
  Bp = std::move(tmp);
}

helib::Ctxt hhmm_sq_shift_compute_A(
    hmm_constant_multipler *cm, const helib::EncryptedArray &ea,
    const std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> &hoisted_A,
    int r) {
  auto *hmmsqcm = static_cast<hhmmsq_constant_multipler *>(cm);
  ASSERT_PRINTF(hmmsqcm != nullptr,
                "Failed to static_cast to `hhmmsq_constant_multipler *`\n");
  const long h_nrow = ea.sizeOfDimension(0);
  auto A_i = hoisted_A->automorph(-r);
  if (r) {
    auto back = hoisted_A->automorph(h_nrow - r);
    hmmsqcm->shift_compute_correct_mask(*A_i, *back, r);
    *A_i += *back;
  }
  return *A_i;
}

void hhmm_sq_bsgs_mt(hmm_metric_collector *collector,
                     hmm_constant_multipler *cm, int threads,
                     const helib::EncryptedArray &ea, helib::Ctxt &AB,
                     const helib::Ctxt &A, const helib::Ctxt &B) {
  auto *hmmsq = static_cast<hmmsq_collector *>(collector);
  ASSERT_PRINTF(hmmsq != nullptr,
                "Failed to static_cast to `hmmsq_collector *`\n");
  auto *hmmsqcm = static_cast<hhmmsq_constant_multipler *>(cm);
  ASSERT_PRINTF(hmmsqcm != nullptr,
                "Failed to static_cast to `hhmmsq_constant_multipler *`\n");
  if (threads == 0) {
    hhmm_sq_bsgs_mt_task(hmmsq, hmmsqcm, ea, AB, A, B);
    return;
  }

  hmmsq->init_noise = std::max(log2_noise_bound(A), log2_noise_bound(B));

  const long h_nrow = ea.sizeOfDimension(0);
  const long h_ncol = ea.sizeOfDimension(1);

  const long f = static_cast<long>(std::ceil(std::sqrt(h_nrow)));
  const long h = (h_nrow + f - 1) / f;

  const helib::PubKey &pk = A.getPubKey();
  helib::Ctxt zero_ctxt(pk);
  zero_ctxt.DummyEncrypt(NTL::ZZX(0));
  helib::Ctxt Ap(zero_ctxt);
  helib::Ctxt Bp(zero_ctxt);
  // helib::Ctxt AB_temp(zero_ctxt);
  std::vector<std::shared_ptr<helib::Ctxt>> C(h_nrow);

  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_A;
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_B;
  // std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_Ap;
  // std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_Bp;

  //! RotateAlign - Start
  std::vector<std::shared_ptr<helib::Ctxt>> A_baby(f);
  std::vector<std::shared_ptr<helib::Ctxt>> B_baby(f);
  std::vector<helib::Ctxt> A_baby_sum(h, zero_ctxt);
  std::vector<helib::Ctxt> B_baby_sum(h, zero_ctxt);
#pragma omp parallel
  {
    INTEL_ITT_TASK_BEGIN(hmmsq_compute_domain, hmmsq_ra_handle); //! ITT - Begin
#pragma omp master               //! No implied barrier
    hmmsq->start_computing_ra(); //! RotateAlign - Start

#pragma omp single nowait
    hoisted_A = std::make_shared<helib::GeneralAutomorphPrecon_FULL>(A, 1, ea);
#pragma omp single nowait
    hoisted_B = std::make_shared<helib::GeneralAutomorphPrecon_FULL>(B, 0, ea);
#pragma omp barrier //! Explicit synchronization

    // Pre-compute ctxts in the baby steps
#pragma omp for
    for (long i = 0; i < f; i++) {
      A_baby[i] = hoisted_A->automorph(-1l * i);
      B_baby[i] = hoisted_B->automorph(-1l * i);
    } //! omp for: Implicit synchronization

    // BSGS
    // FIXME: more parallelism?
#pragma omp for
    for (long i = 0; i < h; i++) {
      // helib::Ctxt A_baby_sum(zero_ctxt);
      // helib::Ctxt B_baby_sum(zero_ctxt);

      const long gaint_step = f * i;

      // Baby Steps
      for (long j = 0; j < f && gaint_step + j < h_nrow; j++) {
        auto A_j = *A_baby[j];
        auto B_j = *B_baby[j];

        hmmsqcm->rotate_align_mask_left(A_j, j + gaint_step);
        hmmsqcm->rotate_align_mask_right(B_j, j + gaint_step);

        // A_baby_sum += A_j;
        // B_baby_sum += B_j;
        A_baby_sum[i] += A_j;
        B_baby_sum[i] += B_j;
      }

      // Gaint Steps
      // ea.rotate1D(A_baby_sum, 1, -1l * gaint_step);
      // ea.rotate1D(B_baby_sum, 0, -1l * gaint_step);
      ea.rotate1D(A_baby_sum[i], 1, -1l * gaint_step);
      ea.rotate1D(B_baby_sum[i], 0, -1l * gaint_step);

      // #pragma omp critical
      //       {
      //         Ap += A_baby_sum;
      //         Bp += B_baby_sum;
      //       }
    } //! omp for: Implicit synchronization

#pragma omp single
    {
      for (int i = 0; i < h; i++) {
        Ap += A_baby_sum[i];
      }
      hmmsq->ra_noise = std::max(log2_noise_bound(Ap), log2_noise_bound(Bp)) -
                        hmmsq->init_noise;

      //! Correct Ap (Correct RotateAlign) - Start
      hmmsq->ra_correct_noise = log2_noise_bound(Ap);
      helib::Ctxt Ap_temp = Ap;
      hmmsqcm->rotate_align_correct_mask(Ap_temp, Ap);
      ea.rotate1D(Ap_temp, 1, -(h_ncol - h_nrow));
      Ap += Ap_temp;
      hmmsq->ra_correct_noise = log2_noise_bound(Ap) - hmmsq->ra_correct_noise;
      //! Correct Ap (Correct RotateAlign) - Completed
      hmmsq->stop_computing_ra(); //! RotateAlign - Completed
      hmmsq->sc_noise = std::max(log2_noise_bound(Ap), log2_noise_bound(Bp));
    } //! omp single: Implicit synchronization
    INTEL_ITT_TASK_END(hmmsq_compute_domain); //! ITT - End

    INTEL_ITT_TASK_BEGIN(hmmsq_compute_domain, hmmsq_sc_handle); //! ITT - Begin
#pragma omp single nowait
    {
      hmmsq->start_computing_sc(); //! ShiftCompute - Start
      hoisted_A =
          std::make_shared<helib::GeneralAutomorphPrecon_FULL>(Ap, 1, ea);
    }
#pragma omp single nowait
    {
      for (int i = 0; i < h; i++) {
        Bp += B_baby_sum[i];
      }
      hoisted_B =
          std::make_shared<helib::GeneralAutomorphPrecon_FULL>(Bp, 0, ea);
    }

#pragma omp barrier //! Explicit synchronization

#pragma omp for
    for (int i = 0; i < h_nrow; i++) {
      // Shift
      auto A_i = hoisted_A->automorph(-i);
      // auto B_i = hoisted_B->automorph(-i);
      C[i] = hoisted_B->automorph(-i);

      // Correct Step: only 'A_i'(i>0) required
      if (i) {
        auto back = hoisted_A->automorph(h_nrow - i);
        hmmsqcm->shift_compute_correct_mask(*A_i, *back, i);
        *A_i += *back;
      }
      // A_i->multiplyBy(*B_i);
      C[i]->multiplyBy(*A_i);

      // #pragma omp critical
      // { AB_temp += *A_i; }
    } //! omp for: Implicit synchronization

#pragma omp master                            //! No implied barrier
    hmmsq->stop_computing_sc();               //! ShiftCompute - Completed
    INTEL_ITT_TASK_END(hmmsq_compute_domain); //! ITT - End
  }
  // AB = std::move(AB_temp);
  AB = std::move(zero_ctxt);
  // NTL::SetNumThreads(omp_get_max_threads());
  for (auto i = 0; i < C.size(); i++) {
    AB += *C[i];
  }
  // NTL::SetNumThreads(1l);
  hmmsq->sc_noise = log2_noise_bound(AB) - hmmsq->sc_noise;
  hmmsq->final_noise = log2_noise_bound(AB);
  hmmsq->left_cap = AB.capacity();
}

void hhmm_sq_bsgs_mt2(hmm_metric_collector *collector,
                      hmm_constant_multipler *cm, int threads,
                      const helib::EncryptedArray &ea, helib::Ctxt &AB,
                      const helib::Ctxt &A, const helib::Ctxt &B) {
  auto *hmmsq = static_cast<hmmsq_collector *>(collector);
  ASSERT_PRINTF(hmmsq != nullptr,
                "Failed to static_cast to `hmmsq_collector *`\n");
  auto *hmmsqcm = static_cast<hhmmsq_constant_multipler *>(cm);
  ASSERT_PRINTF(hmmsqcm != nullptr,
                "Failed to static_cast to `hhmmsq_constant_multipler *`\n");

  const long h_nrow = ea.sizeOfDimension(0);
  const long h_ncol = ea.sizeOfDimension(1);

  const long f = static_cast<long>(std::ceil(std::sqrt(h_nrow)));
  const long h = (h_nrow + f - 1) / f;

  const helib::PubKey &pk = A.getPubKey();
  helib::Ctxt zero_ctxt(pk);
  zero_ctxt.DummyEncrypt(NTL::ZZX(0));
  helib::Ctxt Ap(zero_ctxt);
  helib::Ctxt Bp(zero_ctxt);
  helib::Ctxt AB_temp(zero_ctxt);

  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_A;
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_B;
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_Ap;
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_Bp;

  //! RotateAlign - Start
  hmmsq->start_computing_ra(); //! RotateAlign - Start
  std::vector<std::shared_ptr<helib::Ctxt>> A_baby(f);
  std::vector<std::shared_ptr<helib::Ctxt>> B_baby(f);
  std::vector<helib::Ctxt> A_baby_acc(h_nrow, helib::Ctxt(pk));
  std::vector<helib::Ctxt> B_baby_acc(h_nrow, helib::Ctxt(pk));
  INTEL_ITT_TASK_BEGIN(hmmsq_compute_domain, hmmsq_ra_handle); //! ITT - Begin
#pragma omp parallel
#pragma omp single
  {
#pragma omp task depend(out : hoisted_A)
    hoisted_A = std::make_shared<helib::GeneralAutomorphPrecon_FULL>(A, 1, ea);
    for (long i = 0; i < f; i++) {
#pragma omp task firstprivate(i) depend(in : hoisted_A)
      A_baby[i] = hoisted_A->automorph(-1l * i);
    }
#pragma omp task depend(out : hoisted_B)
    hoisted_B = std::make_shared<helib::GeneralAutomorphPrecon_FULL>(B, 0, ea);
    for (long i = 0; i < f; i++) {
#pragma omp task firstprivate(i) depend(in : hoisted_B)
      B_baby[i] = hoisted_B->automorph(-1l * i);
    }
#pragma omp taskwait

#pragma omp taskloop
    for (long ij = 0; ij < h_nrow; ij++) {
      const auto i = ij / f;
      const auto j = ij % f;
      const auto gaint_step = f * i;
      A_baby_acc[ij] = *A_baby[j];
      B_baby_acc[ij] = *B_baby[j];
      hmmsqcm->rotate_align_mask_left(A_baby_acc[ij], j + gaint_step);
      hmmsqcm->rotate_align_mask_right(B_baby_acc[ij], j + gaint_step);
    } //! omp for: Implicit synchronization

#pragma omp taskloop
    for (long i = 0; i < h; i++) {
      helib::Ctxt A_baby_sum(zero_ctxt);
      helib::Ctxt B_baby_sum(zero_ctxt);
      const auto gaint_step = f * i;

      // Baby Steps
      for (long j = 0; j < f && gaint_step + j < h_nrow; j++) {
        const auto idx = gaint_step + j;
        A_baby_sum += A_baby_acc[gaint_step + j];
        B_baby_sum += B_baby_acc[gaint_step + j];
      }

      // Gaint Steps
      ea.rotate1D(A_baby_sum, 1, -1l * gaint_step);
      ea.rotate1D(B_baby_sum, 0, -1l * gaint_step);

#pragma omp critical
      {
        Ap += A_baby_sum;
        Bp += B_baby_sum;
      }
    }

#pragma omp task depend(out : Ap)
    {
      //! Correct Ap (Correct RotateAlign) - Start
      helib::Ctxt Ap_temp = Ap;
      hmmsqcm->rotate_align_correct_mask(Ap_temp, Ap);
      ea.rotate1D(Ap_temp, 1, -(h_ncol - h_nrow));
      Ap += Ap_temp;
      //! Correct Ap (Correct RotateAlign) - Completed
      hmmsq->stop_computing_ra(); //! RotateAlign - Completed
    }
    INTEL_ITT_TASK_END(hmmsq_compute_domain); //! ITT - End

    INTEL_ITT_TASK_BEGIN(hmmsq_compute_domain, hmmsq_sc_handle); //! ITT - Begin
    hmmsq->start_computing_sc(); //! ShiftCompute - Start
#pragma omp task depend(in : Ap)
    hoisted_Ap =
        std::make_shared<helib::GeneralAutomorphPrecon_FULL>(Ap, 1, ea);
#pragma omp task
    hoisted_Bp =
        std::make_shared<helib::GeneralAutomorphPrecon_FULL>(Bp, 0, ea);
#pragma omp taskwait

#pragma omp taskloop
    for (int i = 0; i < h_nrow; i++) {
      // Shift
      auto A_i = hoisted_Ap->automorph(-i);
      auto B_i = hoisted_Bp->automorph(-i);

      // Correct Step: only 'A_i'(i>0) required
      if (i) {
        auto back = hoisted_Ap->automorph(h_nrow - i);
        hmmsqcm->shift_compute_correct_mask(*A_i, *back, i);
        *A_i += *back;
      }
      A_i->multiplyBy(*B_i);

#pragma omp critical
      { AB_temp += *A_i; }
    }
  }
  INTEL_ITT_TASK_END(hmmsq_compute_domain); //! ITT - End
  hmmsq->stop_computing_sc();               //! ShiftCompute - Completed
  AB = std::move(AB_temp);
}